package com.ckjava.xutils.access;

import com.ckjava.xutils.aop.AopBase;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

/**
 * qps 访问控制工具类
 */
public class IpControlUtils {

    private static final Logger logger = LoggerFactory.getLogger(AopBase.class);

    /**
     * 单位时间 1000 毫秒
     */
    private static final int timePer = 1000;

    /**
     * 单位时间最大访问次数
     */
    private static final int maxPer = 10;

    /**
     * 单个 ip 最长有效时间
     */
    private static final int maxAlive = 10000;

    /**
     * 记录每个IP 对应的 访问控制实体：IpControl
     */
    public static ConcurrentHashMap<String, IpControl> controlConcurrentHashMap = new ConcurrentHashMap<>();

    /**
     * 定时清理合法的 ip
     */
    public static ScheduledExecutorService scheduledExecutorService = Executors.newScheduledThreadPool(1);

    /**
     * 系统启动的时候自动启动 定时清理合法的 ip 的线程
     */
    static {
        scheduledExecutorService.scheduleAtFixedRate(() -> {
            try {
                // 1. 找到存活 10 秒以上的
                List<String> removeList = new ArrayList<>();
                controlConcurrentHashMap.forEach((t,u) -> {
                    String ip = t;
                    IpControl ipControl = u;
                    long timePeriod = System.currentTimeMillis() - ipControl.getFirstAccessTime().get();
                    if (timePeriod > maxAlive) { // 将存活超过最大时间的数据删除
                        removeList.add(ip);
                    }
                });

                // 2. 通过 key 进行删除
                Optional.of(removeList)
                        .filter(t -> t.size()>0)
                        .ifPresent(t -> {
                            logger.info("removeList size:{}", t.size());
                            t.forEach(s -> controlConcurrentHashMap.remove(s));
                        });
            } catch (Throwable e) {
                logger.error("scheduleAtFixedRate has error", e);
            }

        }, 10, 10, TimeUnit.SECONDS);
    }

    /**
     * 判断 ip 是否符合访问控制策略
     * @param ip ip 地址
     * @return
     */
    public static boolean isAccess(String ip) {
        IpControl ipControl = controlConcurrentHashMap.get(ip);
        if (ipControl == null) { // 首次访问
            initIpControl(ip);
            return true;
        } else {
            long timePeriod = System.currentTimeMillis() - ipControl.getFirstAccessTime().get();
            if (timePeriod <= 0) {
                return true;
            }
            if (ipControl.getCount().get() >= maxPer) { // 访问次数超过 限制
                if (timePeriod < timePer) { // 不到1秒，不通过并将首次访问时间加上 100 毫秒
                    ipControl.getFirstAccessTime().getAndAdd(100);
                    return false;
                } else {  // 超过1秒, 通过
                    initIpControl(ip);
                    return true;
                }
            } else { // 访问次数不超过 限制
                if (timePeriod < timePer) { // 不到1秒, 通过
                    ipControl.getCount().incrementAndGet();
                    return true;
                } else { // 大于1秒, 通过
                    initIpControl(ip);
                    return true;
                }
            }
        }
    }

    /**
     * 初始化 ip 访问策略
     *
     * @param ip
     */
    private static void initIpControl(String ip) {
        IpControl ipControl = new IpControl();
        ipControl.setCount(new AtomicLong(1));
        ipControl.setFirstAccessTime(new AtomicLong(System.currentTimeMillis()));
        controlConcurrentHashMap.put(ip, ipControl);
    }

    /**
     * ip 访问控制实体
     */
    private static class IpControl {

        /**
         * ip 首次访问时间
         */
        private AtomicLong firstAccessTime;

        /**
         * 记录 firstAccessTime 后的访问量
         */
        private AtomicLong count;

        public AtomicLong getFirstAccessTime() {
            return firstAccessTime;
        }

        public void setFirstAccessTime(AtomicLong firstAccessTime) {
            this.firstAccessTime = firstAccessTime;
        }

        public AtomicLong getCount() {
            return count;
        }

        public void setCount(AtomicLong count) {
            this.count = count;
        }

        public IpControl(AtomicLong firstAccessTime, AtomicLong count) {
            this.firstAccessTime = firstAccessTime;
            this.count = count;
        }

        public IpControl() {
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) return true;
            if (o == null || getClass() != o.getClass()) return false;
            IpControl ipControl = (IpControl) o;
            return Objects.equals(firstAccessTime, ipControl.firstAccessTime) &&
                    Objects.equals(count, ipControl.count);
        }

        @Override
        public int hashCode() {
            return Objects.hash(firstAccessTime, count);
        }
    }
}
